from werkzeug.datastructures import FileStorage

from owl_admin.ext import db
from owl_ai.domain.knowledge.knowledge_service import KnowledgeFileService, KnowledgeService, VectorStoreService
from owl_common.sqlalchemy.transaction import Transactional


class KnowledgeBaseAppService:

    @classmethod
    @Transactional(db.session)
    def upload(cls, kid, file: FileStorage):
        # 上传文件到磁盘中
        disk_name, file_path = KnowledgeFileService.upload_to_disk(file)

        knowledge_file_record = KnowledgeFileService.save_db_record(kid, disk_name, file)

        knowledge_base = KnowledgeService.find_by_id(kid)

        # 加载文件
        documents = KnowledgeFileService.load_file(file_path)

        # 文件切分
        split_documents = KnowledgeFileService.split_file(documents, knowledge_base.chunk_strategy)

        # 文档向量化并入库
        ids = VectorStoreService.embedding_docs_save(knowledge_base, knowledge_file_record, docs=split_documents)
        return ids

    @classmethod
    @Transactional(db.session)
    def knowledge_file_list(cls, kid):
        return KnowledgeFileService.list(kid)

    @classmethod
    @Transactional(db.session)
    def similarity_search(cls, kid, query, top_k):
        knowledge_base = KnowledgeService.find_by_id(kid)
        docs = VectorStoreService.similarity_search(knowledge_base, query, top_k)
        return docs
